#include "common.h"
#include "ms_model.h"
#include "napi/native_api.h"
#include "torch_fft_ifft.h"
#include "video_processor.h"
#include "wav_header.h"
#include <dlfcn.h>
#include <hilog/log.h>
#include <iostream>
#include <mindspore/context.h>
#include <mindspore/model.h>
#include <mindspore/status.h>
#include <mindspore/tensor.h>
#include <mindspore/types.h>
#include <random>
#include <stdio.h>
#include <string>

#include <filesystem>
#include <iostream>

namespace fs = std::__fs::filesystem;
int Process(OH_AI_ModelHandle handle_) {
    std::random_device rd;                                 // 获取真随机种子（硬件熵源）
    std::mt19937 gen(rd());                                // 初始化梅森旋转引擎（高质量伪随机）
    std::uniform_real_distribution<double> dist(0.0, 1.0); // [0.0, 1.0)区间
    int random_num = dist(gen);                            // 生成随机数
    auto inputs = OH_AI_ModelGetInputs(handle_);
    auto outputs = OH_AI_ModelGetOutputs(handle_);
    for (int c = 0; c < inputs.handle_num; ++c) {
        float *data = (float *)OH_AI_TensorGetMutableData(inputs.handle_list[c]);
        size_t num = OH_AI_TensorGetElementNum(inputs.handle_list[c]);
        for (int i = 0; i < num; ++i) {
            data[i] = float(dist(gen));
        }
    }

    LOGI("predict start ...\n");
    auto start = std::chrono::high_resolution_clock::now();
    OH_AI_ModelPredict(handle_, inputs, &outputs, nullptr, nullptr);
    auto end = std::chrono::high_resolution_clock::now();
    LOGI("predict end ...\n");
    return std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
}


static napi_value InitEnv(napi_env env, napi_callback_info info) {
    // 传入数据处理
    size_t argc = 1;
    napi_value argv[1] = {nullptr};
    napi_get_cb_info(env, info, &argc, argv, nullptr, nullptr);
    char model_name[128];
    size_t name_size = 128;
    napi_get_value_string_utf8(env, argv[0], model_name, 128, &name_size);
    LOGI("PATH： %{public}s", model_name);

//    LOGI("model input size: %{public}d", input.size());
//    LOGI("model input[0] size: %{public}d", input[0].size());
//    LOGI("model input[0][0] size: %{public}d", input[0][0].size());
//    for(int i = 0; i < 100; ++i) {
//        LOGI("%{public}f, %{public}f", input[0][i][0], input[0][i][1]);
//    }
    return nullptr;
}

void print_directory_contents(const fs::path &path) {
    if (!fs::exists(path) || !fs::is_directory(path)) {
        LOGE("错误：路径不存在或不是目录");
        return;
    }

    try {
        // 使用 directory_iterator 替代 recursive_directory_iterator 进行非递归遍历
        for (const auto &entry : fs::directory_iterator(path)) {
            if (entry.is_directory()) {
                LOGI("[目录] %{public}s", entry.path().string().c_str());
            } else if (entry.is_regular_file()) {
                uintmax_t file_size = entry.file_size();
                LOGI("[文件] %{public}s (%{public}lu 字节)", entry.path().filename().string().c_str(), file_size);
            } else {
                LOGI("[其他] %{public}s", entry.path().filename().string().c_str());
            }
        }
    } catch (const fs::filesystem_error &ex) {
        LOGE("文件系统错误: %{public}s", ex.what());
    }
}

std::vector<std::vector<std::vector<float>>> get_first_time_frame(
    const std::vector<std::vector<std::vector<float>>>& stft_result) {
    
    int num_freq_bins = stft_result.size();
    if (num_freq_bins == 0) {
        return {};
    }
    
    // 创建结果向量，维度为 [num_freq_bins, 1, 2]
    std::vector<std::vector<std::vector<float>>> first_frame(
        num_freq_bins, 
        std::vector<std::vector<float>>(
            1,  // 只有一个时间帧
            std::vector<float>(2, 0.0f)  // 实部和虚部
        )
    );
    
    // 复制第一个时间帧的数据
    for (int freq_bin = 0; freq_bin < num_freq_bins; ++freq_bin) {
        // 假设每个频率bin至少有一个时间帧
        if (!stft_result[freq_bin].empty()) {
            first_frame[freq_bin][0][0] = stft_result[freq_bin][0][0]; // 实部
            first_frame[freq_bin][0][1] = stft_result[freq_bin][0][1]; // 虚部
        }
    }
    
    return first_frame;
}
std::vector<std::vector<std::vector<float>>> create_tensor_257_1_2() {
    std::vector<std::vector<std::vector<float>>> tensor(
        257, 
        std::vector<std::vector<float>>(
            1, 
            std::vector<float>(2, 0.5f)
        )
    );
    return tensor;
}
static napi_value Run(napi_env env, napi_callback_info info) {
    VideoProcessor videoProcessor;
    int out_sample_rate;
    const std::string wav_file("/data/storage/el1/bundle/entry/resources/resfile/mix.wav");
    const std::string model_name("/data/storage/el1/bundle/entry/resources/resfile/gtcrn_model_slim.ms");


    auto start = std::chrono::high_resolution_clock::now();
    const std::vector<float> mix = videoProcessor.read_wav_file(wav_file, out_sample_rate);
//    for (int i = 0; i < 10; ++i) {
//        LOGI("mix %{public}d: %{public}f", i, mix[i]);
//    }
    //   LOGI("mix len is %{public}d, sf is %{public}d", mix.size(), out_sample_rate);

    auto input = stft(mix, out_sample_rate);
    for (int i = 0; i < 10; ++i) {
        LOGI("%{public}f, %{public}f", input[0][i][0], input[0][i][1]);
    }
    // 正确获取维度信息
    int batch_size = input.size();            // 应该是1
    int freq_bins = input[0].size();          // 应该是257 (512/2+1)
    int time_frames = input[0][0].size() / 2; // 每个时间帧有实部和虚部

    LOGI("STFT维度: batch=%d, freq=%d, time=%d", batch_size, freq_bins, time_frames);
    
    MSModel model;
    model.CreateMsLiteModelCPUAsync(model_name);
     //auto input1 =create_tensor_257_1_2();
    model.FillInputTensor(0, input);
    LOGI("predict start ...\n");
    const float *out = model.Predict();
    LOGI("-------");
    for(int i =0; i < 10; ++i) {
        LOGI("%{public}f", out[i]);
    }
    LOGI("-------");
    // 使用正确的维度进行转换
    auto ad = float_array_to_vector3d(out, batch_size, freq_bins, time_frames);
//    for (int i = 0; i < 10; ++i) {
//        LOGI("%{public}f, %{public}f", ad[0][i][0], ad[0][i][1]);
//    }
    auto out1 = istft(ad);
//    for (int i = 0; i < 10; ++i) {
//        LOGI("out %{public}d: %{public}f", i, out1[i]);
//    }
    auto end = std::chrono::high_resolution_clock::now();
    auto a = std::chrono::duration_cast<std::chrono::milliseconds>(end - start).count();
    // 保存原始输入音频（调试用）
    SaveAudioToWav(mix.data(), mix.size(), out_sample_rate, 1, "/data/storage/el2/base/haps/raw_audio.wav");

    // 保存处理后的音频
    SaveModelOutputToWav(out1.data(), out1.size(), out_sample_rate, "/data/storage/el2/base/haps/enhanced_audio.wav");

    print_directory_contents("/data/storage/el2/base/haps/");

    napi_value res;
    napi_create_uint32(env, a, &res);
    return res;
}

static napi_value SwitchPredictContext(napi_env env, napi_callback_info info) {
    // 传入数据处理
    size_t argc = 2;
    napi_value argv[2] = {nullptr};
    napi_get_cb_info(env, info, &argc, argv, nullptr, nullptr);
    char model_name[128];
    size_t name_size = 128;
    napi_get_value_string_utf8(env, argv[0], model_name, 128, &name_size);

    return nullptr;
}


EXTERN_C_START
static napi_value Init(napi_env env, napi_value exports) {
    napi_property_descriptor desc[] = {
        {"run", nullptr, Run, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"initEnv", nullptr, InitEnv, nullptr, nullptr, nullptr, napi_default, nullptr},
        {"switchPredictContext", nullptr, SwitchPredictContext, nullptr, nullptr, nullptr, napi_default, nullptr}};
    napi_define_properties(env, exports, sizeof(desc) / sizeof(desc[0]), desc);
    return exports;
}
EXTERN_C_END

static napi_module demoModule = {
    .nm_version = 1,
    .nm_flags = 0,
    .nm_filename = nullptr,
    .nm_register_func = Init,
    .nm_modname = "entry",
    .nm_priv = ((void *)0),
    .reserved = {0},
};

extern "C" __attribute__((constructor)) void RegisterEntryModule(void) { napi_module_register(&demoModule); }
